import Core.MySQLDB as DataBase
import Core.Gadget as Gadget
import mysql.connector
from datetime import date, datetime, timedelta
from copy import copy
from Core.DataSeries import DataSeries
import pandas as pd

class JDMySQLDB(DataBase.MySQLDB):

    def __init__(self, address, port, config):
        self.name = "JDMySQLDB 单密码数据库"
        print("建立数据库对象", self.name)
        # print("Connecting MySQL")
        self.address = address
        self.port = port
        self.config = config
        self.contexts = {}
        # print("Connected MySQL")

    def _Connect(self, databaseName):
        #
        config = self.config.get(databaseName)
        if not config:
            print("Invalid Database Name", databaseName)
            return

        #
        host = None
        if "Address" in config:
            host = config.get("Address")
        else:
            host = self.address
        #
        port = None
        if "Port" in config:
            port = config.get("Port")
        else:
            port = self.port
        #
        context = mysql.connector.connect(user=config["Username"],
                                          password=config["Password"],
                                          db=databaseName,
                                          host=host,
                                          port=port,
                                          buffered=True) # 占用大量内存，但是避免handle_unread_result错误
        #
        self.contexts[databaseName] = context
        return context

    # 为了兼容旧代码，需要重定位数据库和表名
    def _ReDirectDatabaseAndTable(self, databaseName, tableName):
        #
        if databaseName.lower() == "dailybar":
            return "financial_data", tableName + "_dailybar"
        #
        if databaseName.lower() == "instruments":
            if tableName == None or tableName == "":
                return "financial_data", "Instrument"
            return "financial_data",  "Instrument_" + tableName
        #
        if databaseName.lower() == "stock":
            return "financial_data", "Stock_" + tableName
        #
        if databaseName.lower() == "mutualfund":
            return "financial_data", "mutualfund_" + tableName
        #
        if databaseName.lower() == "bond":
            return "financial_data", "bond_" + tableName
        #
        if databaseName.lower() == "future":
            return "financial_data", "future_" + tableName
        #
        if databaseName.lower() == "index":
            return "financial_data", "index_" + tableName
        #
        if databaseName.lower() == "option":
            return "financial_data", "option_" + tableName
        #
        return databaseName, tableName

    # JD数据库只能接受小写表头
    def _MakeLower(self):
        pass

    def Find(self, databaseName, tableName, filter={}, sort=[], projection={}):
        databaseName, tableName = self._ReDirectDatabaseAndTable(databaseName, tableName)
        return super().Find(databaseName, tableName, filter, sort, projection)

    def Insert(self, databaseName, tableName, document):
        databaseName, tableName = self._ReDirectDatabaseAndTable(databaseName, tableName)
        return super().Insert(databaseName, tableName, document)

    def Upsert(self, databaseName, tableName, target, document):
        databaseName, tableName = self._ReDirectDatabaseAndTable(databaseName, tableName)
        return super().Upsert(databaseName, tableName, target, document)

    def Update(self, databaseName, tableName, target, document):
        databaseName, tableName = self._ReDirectDatabaseAndTable(databaseName, tableName)
        return super().Update(databaseName, tableName, target, document)

    def Upsert_Many(self, databaseName, tableName, targets, documents):
        databaseName, tableName = self._ReDirectDatabaseAndTable(databaseName, tableName)
        return super().Upsert_Many(databaseName, tableName, targets, documents)

    def Insert_Many(self, databaseName, tableName, documents):
        databaseName, tableName = self._ReDirectDatabaseAndTable(databaseName, tableName)
        return super().Insert_Many(databaseName, tableName, documents)

    def GetDataFrame(self, databaseName, tableName, filter=[], projection=[], datetime1=None, datetime2=None, sort=None, date_field="datetime"):
        databaseName, tableName = self._ReDirectDatabaseAndTable(databaseName, tableName)
        return super().GetDataFrame(databaseName, tableName, filter, projection, datetime1, datetime2, sort, date_field)

    def GetDataFrame2(self, symbol, dataType=None, instrumentType=None, datetime1=None, datetime2=None):
        #
        databaseName = "DailyBar"
        tableName = instrumentType
        #
        if dataType == "DailyBar":
            databaseName = "financial_data"
            tableName = instrumentType + "_dailybar"
        #
        sort = [("DateTime", 1)]
        documents = self.Find(databaseName, tableName, filter={"Symbol": symbol}, sort=sort)
        df = Gadget.DocumentsToDataFrame(documents)

        # 必要的大小写转换
        convertUpper = {"close":"Close", "bclose":"BClose", "symbol":"Symbol", "date":"Date", "datetime":"DateTime", "adjfactor":"AdjFactor"}
        realConvertUpper = {}
        for key, value in convertUpper.items():
            if key in df.columns:
                realConvertUpper[key] = value
        #
        df.rename(columns=realConvertUpper, inplace=True)

        # ---convert datetime---
        convertDateTimeFields = ["Date", "Datetime"]
        for field in convertDateTimeFields:
            if field in df.columns:
                df[field] = pd.to_datetime(df[field])
        #
        if datetime1 and datetime2:
            df = df[(df["DateTime"] >= datetime1) & (df["DateTime"] <= datetime2)]
        elif datetime1:
            df = df[df["DateTime"] >= datetime1]
        elif datetime2:
            df = df[df["DateTime"] <= datetime2]

        # 引入新数据库时遭遇的问题
        convertFloatFields = ["Close", "BClose"]
        for field in convertFloatFields:
            if field in df.columns:
                df[field] = df[field].astype("float")
        #
        return df

    def SaveDataSeries(self, dataSeries, dataType=None, instrumentType=None, isUpSert=True):
        #
        documents = []
        for item in dataSeries.items:
            document = copy(item)
            #
            if "DataSeries" in item:
                del document["DataSeries"]
            #
            documents.append(document)

        #
        if dataType.lower() == "dailybar":
            databaseName = "financial_data"
            tableName = instrumentType + "_dailybar"
        #
        elif dataType.lower() == "factor":
            databaseName = "factor"
            tableName = dataSeries.model_name
        else:
            return
        #
        self.Upsert_Many(databaseName, tableName, {}, documents)
        #
        print(str(datetime.now()), "Saved DataSeries:", dataSeries.model_name, "#Elements", str(dataSeries.Count()))